package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{IndexToStringEncoder, IndexToStringEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author tangjinyuan
  * @Date 2019/3/27 14:32
  * @Description
  * @Reviewer
  */
class IndexToStringEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "categoryIndex",
      |	"outputCol": "originalCategory",
      |	"labels": ["a","b","c"]
      |}
    """.stripMargin

  val sie = new IndexToStringEncoder()
  val param: IndexToStringEncoderParam = sie.parseParam(new ToolClass().encrypt(paramStr))

  test("sie parameter") {
    assert("categoryIndex".equals(param.inputCol) && "originalCategory".equals(param.outputCol) && param.labels.length==3)
  }


  test("sie transformation") {
    var data = TestingDataGenerator.spark.createDataFrame(
      Seq((0, 0.0), (1, 2.0), (2, 1.0), (3, 0.0), (4, 0.0), (5, 1.0))
    ).toDF("id", "categoryIndex")

    var data2 = TestingDataGenerator.spark.createDataFrame(
      Seq(( 0, 0.0), (1, 2.0), (2, 1.0))
    ).toDF("id", "categoryIndex")
    val res = sie.train(param, data).transform(data2)

    // 输出结果 行数为3个
    res.show()
    assert(res.select(param.outputCol).rdd.count() == 3)
  }

}
